import random
from .models import LogistQuestion


def topic_sampling(
    topic_ratios,
    total
):
    """
    按主题比例从 LogistQuestion 中随机抽取 total 道题
    topic_ratios: {topic: weight}
    返回题目对象列表
    """
    sum_t = sum(topic_ratios.values())
    plan = {}
    base = total / sum_t
    for t, wt in topic_ratios.items():
        cnt = int(wt * base)
        plan[t] = cnt
    # 调整剩余
    assigned = sum(plan.values())
    for key in sorted(plan.keys(), key=lambda x: plan[x] - int(plan[x]), reverse=True):
        if assigned < total:
            plan[key] += 1
            assigned += 1
    selected = []
    for topic, cnt in plan.items():
        if cnt <= 0:
            continue
        ids = list(LogistQuestion.objects.filter(
            topic=topic
        ).values_list('id', flat=True))
        if cnt > len(ids):
            raise ValueError(f"[{topic}] 题目不足: 需{cnt}, 有{len(ids)}")
        sampled = random.sample(ids, cnt)
        selected.extend(list(LogistQuestion.objects.filter(id__in=sampled)))
    random.shuffle(selected)
    return selected